import json
import matplotlib.pyplot as plt

def plot_loss_curve(jsonl_path, tar_path):
    steps = []
    losses = []
    # 读取jsonl
    with open(jsonl_path, 'r', encoding='utf-8') as f:
        for line in f:
            if not line.strip():
                continue
            obj = json.loads(line)
            steps.append(obj.get('current_steps', obj.get('step', None)))
            losses.append(obj.get('loss', None))
    
    # 绘图
    plt.figure(figsize=(8,4), dpi=60)
    plt.plot(steps, losses, marker='.', linestyle='-', alpha=0.8)
    plt.xlabel("Step")
    plt.ylabel("Loss")
    plt.title("Training Loss Curve")
    plt.grid(True, alpha=0.2)
    plt.tight_layout()
    plt.savefig(tar_path, dpi=60)
    plt.close()

# 你可以按如下调用：
plot_loss_curve("./experiments/saves/pub_med/lora/trainer_log.jsonl", "output.png")

